from torch.utils.data import DataLoader, Dataset

import torch


class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


data = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
custom_dataset = CustomDataset(data)

dataloader = DataLoader(custom_dataset, batch_size=2, shuffle=False)

for batch in dataloader:
    print(batch)


